import numpy as np
import sys
import os
import json
from skimage import io
import cv2
from scipy.io import savemat

# sys.path.append('/home/lab/wjSun/caffe-master/python')
sys.path.insert(0, '/home/dell/caffe-master/python')
import caffe
from fine_tune_pretraind_keras_vgg_jiao_dong import read_and_normalize_test_data, merge_several_folds_mean, merge_several_folds_geom,create_submission

caffe.set_device(1)  # if we have multiple GPUs, pick the first one
caffe.set_mode_gpu()

#pretrained_root = '/media/lab/labdisk/home/lab1314/xryu/pretrined_model/caffe/ResNet/'
#model_def = pretrained_root + 'ResNet-50-deploy.prototxt'
#model_weights = pretrained_root + 'ResNet-50-model.caffemodel'

model_def = '/home/dell/Rui/kaggle_ddd_finetune_caffe/fine_tune_googlenet/googlenet_test.prototxt'
model_weights = '/home/dell/Rui/kaggle_ddd_finetune_caffe/fine_tune_googlenet/fine_tuned/_iter_426000.caffemodel'
#126000 top1 acc=0.90, top5 acc=0.98
#156000 top1 0.9015
#426000 top1 0.91
#solver_path = '/home/lab/xryu/code/DistractedDriverDetection/kaggle_ddd_finetune_caffe/fine_tune_googlenet/solver_multistep.prototxt'
imagenet_mean_path = '/home/dell/Rui/kaggle_ddd_finetune_caffe/fine_tune_googlenet/fine_tuned/imagenet_mean.npy'
IMAGE_FILE = '/media/lab/labdisk/home/lab1314/xryu/data/DistractedDriverDetection/imgs/test/img_24.jpg'
TEST_IMG_ROOT='/media/dell/delldisk/dell/Rui/data/Kaggle/DistractedDriverDetection/imgs/test'



net = caffe.Net(model_def,  # defines the structure of the model
                       model_weights,  # contains the trained weights
                       caffe.TEST
                       )  # use test mode (e.g., don't perform dropout)
batch_size=32
net.blobs['data'].reshape(batch_size,3,224,224)

imagenet_mean = np.load(imagenet_mean_path)
print imagenet_mean.shape

test_data, test_id = read_and_normalize_test_data(224,224,3)
print test_data.shape

n_epoch=len(test_data)/batch_size
yfull_test=[]
for e in range(n_epoch):
    net.blobs['data'].data[...]=test_data[e*batch_size:(e+1)*batch_size]
    output=net.forward()
    prediction1=output['prob1']
    prediction2=output['prob2']
    prediction3=output['prob3']
    yfull_test.append(prediction3)

info_string='fine_tune_googlenet_multistep_solver'
test_res = merge_several_folds_mean(yfull_test, nfolds=1)
create_submission(test_res, test_id, info_string)

'''
img=cv2.imread(IMAGE_FILE)
img=cv2.resize(img,(224,224))
#img_arr=cv2.transpose(img,(2,0,1))
img_arr=np.asarray(img)
img_arr=img_arr.reshape((1,3,224,224))
#img_arr=img_arr.transpose((0,3,1,2))

#solver = caffe.get_solver(solver_path)
#net.set_raw_scale('data', 255)
#net.set_mean('data', np.load(imagenet_mean))
#input_image = caffe.io.load_image(IMAGE_FILE)

net.blobs['data'].data[...]=img_arr[:]
output=net.forward()
prediction1=output['prob1']
prediction2=output['prob2']
prediction3=output['prob3']
#prediction = net.predict([input_image])
print 'prediction shape:', prediction1[0].shape
print 'predicted class:', prediction1[0].argmax()
print prediction1

print 'prediction shape:', prediction2[0].shape
print 'predicted class:', prediction2[0].argmax()
print prediction2

print 'prediction shape:', prediction3[0].shape
print 'predicted class:', prediction3[0].argmax()
print prediction3

#[[ 0.00211135  0.00643646  0.00143373  0.00237239  0.00546147  0.00707581
   #0.00302396  0.0803294   0.01793006  0.87382537]]
'''
